//******************************************************************************
// cmix_3track.cc - A BASF module for reading MDST files and extracting 
// quantities of interest for events consistent with a mu+mu-pi0pi0 final state.
//
// Date: June 2012
// Author: craigb (craig dot bookwalter at pnnl dot gov)
//******************************************************************************

// Standard headers
#include <algorithm>
#include <exception>
#include <iostream>
#include <fstream>
#include <set>
#include <stdexcept>
#include <string>
#include <sstream>
#include <utility>
#include <vector>

// Belle-specific headers
#include "particle/Particle.h"
#include "particle/PID.h"
#include "particle/utility.h"
#include "event/BelleEvent.h"
#include "basf/module.h"
#include "basf/module_descr.h"
#include "benergy/BeamEnergy.h"
#include "kid/atc_pid.h"
#include "ip/IpProfile.h"
#include "helix/Helix.h"
#include "kfitter/kvertexfitter.h"
#include "eid/eid.h"

#include <mdst/mdst.h>

#include "panther/panther.h"
#include MDST_H
#include MDST_OBS_H
#include EVTCLS_H
#include BELLETDF_H
#include RUN_INFO_H

// Headers written by me
#include "rows/cmix_countrow.hpp"
#include "rows/cmix_chargedrow.hpp"

#include <pytpp.hpp>
#include <csbbelle.hpp>
#include <csbbelle_runinforow.hpp>

using std::string;
using std::vector;
using std::cout;
using std::cerr;
using std::endl;

static const double RADS_TO_DEG = 360./(2*3.14159265359);
static const double D0_MASS= 1.86486;
static const double D0_SIGMA=0.0062;


#ifdef BELLE_NAMESPACE
namespace Belle {
#endif //BELLE_NAMESPACE

  // All BASF modules inherit from the "Module" class, defined in `basf/module.h'
  // They are written as a single class. The BASF executable loads this class
  // (compiled into a shared object) dynamically. The module constructor, `init' 
  // and `hist_def' methods are all called before reading events from the event 
  // source. The first event is read and BASF calls `begin_run()'. Every time 
  // the experimental run number (contained in the datastream) changes, `end_run()' 
  // will be called and then `begin_run()' again.
  //
  // The `event()' function does the event-by-event processing; in the case of 
  // this module, saving some quantities of interest to an HDF5 dataset for 
  // later study. 
  //
  // When the event source is exhausted, the term() function is called, followed 
  // by the class destructor.
  //----------------------------------------------------------------------------
  class cmix_3track: public Module {
  public:
    cmix_3track();
    virtual ~cmix_3track();
    void init(int *);
    void term();
    void disp_stat(const char*);
    void hist_def();
    void event(BelleEvent*, int*);
    void begin_run(BelleEvent*, int*);
    void end_run(BelleEvent*, int*);

    char outpath_stub[4096];
    char inpath[4096]; 

    static HepPoint3D IP;
    static HepSymMatrix IPerr;

  private:
    bool _write_h5;
    bool _write_hbook;

    int _nevents_tot;
    int _combos_found;

    // D0*+ candidates
    pyt::table_writer<cmix_chargedrow_t>* _ntpos; 
    cmix_chargedrow_t* _rowpos;

    // D0*- candidates
    pyt::table_writer<cmix_chargedrow_t>* _ntneg; 
    cmix_chargedrow_t* _rowneg;

    pyt::table_writer<cmix_countrow_t>* _countnt; 
    cmix_countrow_t* _countrow;

    pyt::table_writer<csbbelle_runinforow_t>* _runnt; 
    csbbelle_runinforow_t* _runrow;
  };

  HepPoint3D cmix_3track::IP = HepPoint3D(0,0,0);
  HepSymMatrix cmix_3track::IPerr = HepSymMatrix(3,0);

  //----------------------------------------------------------------------------
  // This function registers the module with the BASF executeable.
  // Any parameters defined in the control script (`cmix_3track.basf') are
  // declared and matched with public member variables of the module class.
  //----------------------------------------------------------------------------
  extern "C" Module_descr *mdcl_cmix_3track()
  {

    cmix_3track *tmp = new cmix_3track();
    Module_descr *dscr = new Module_descr ( "cmix_3track", tmp);
    dscr->define_param("outpath_stub", "Path stub to output file", "S", 4096, tmp->outpath_stub);
    dscr->define_param("inpath", "Input file", "S", 4096, tmp->inpath);

    BeamEnergy::define_global(dscr);
    IpProfile::define_global(dscr);

    return dscr;
  }
  
  //----------------------------------------------------------------------------
  // Default constructor
  //----------------------------------------------------------------------------
  cmix_3track::cmix_3track() : 
  _combos_found(0), 
    _rowpos(new cmix_chargedrow_t),
    _rowneg(new cmix_chargedrow_t),
    _countrow(new cmix_countrow_t),
    _runrow(new csbbelle_runinforow_t),
    _ntpos(NULL),
    _ntneg(NULL),
    _countnt(NULL),
    _runnt(NULL)
  {
    // Initialize member vars
  }

  //----------------------------------------------------------------------------
  // Default destructor
  //----------------------------------------------------------------------------  
  cmix_3track::~cmix_3track() {
    delete _ntpos;
    delete _ntneg;

    delete _rowpos;
    delete _rowneg;

    delete _countnt;
    delete _countrow;

    delete _runnt;
    delete _runrow;
  }

  //----------------------------------------------------------------------------
  // In init(), I define one or more output files based on the values of the 
  // parameters passed in from the control scripts.
  //----------------------------------------------------------------------------    
  void cmix_3track::init(int* status) {
    string path_stub(outpath_stub);


    string outpath = path_stub + ".pos.h5";      
    string title = "Charm-mixing analysis with D0*+ candidates";
    _ntpos = new pyt::table_writer<cmix_chargedrow_t>(outpath, "nt", title,
						      0, 1);
    _ntpos->table_setattr("mdst_path", inpath);
    _ntpos->table_setattr("module", "cmix");


    outpath = path_stub + ".neg.h5";      
    title = "Charm-mixing analysis with D0*- candidates";
    _ntneg = new pyt::table_writer<cmix_chargedrow_t>(outpath, "nt", title, 
						      0, 1);
						  
    _ntneg->table_setattr("mdst_path", inpath);
    _ntneg->table_setattr("module", "cmix");
    

    outpath = path_stub + ".count.h5";
    title = "Event-counting ntuple for charm-mixing analysis";
    _countnt = new pyt::table_writer<cmix_countrow_t>(outpath, "nt", title, 
						      0, 1);
    _countnt->table_setattr("mdst_path", inpath);
    _countnt->table_setattr("module", "cmix");

    outpath = path_stub + ".runinfo.h5";
    title = "Belle run information";
    _runnt = new pyt::table_writer<csbbelle_runinforow_t>(outpath, "nt", title,
							  0, 0);
							

    cout << "CRAIGB : Analyzing file " << inpath << endl;
  }
  
  //----------------------------------------------------------------------------
  // hist_def() I leave empty
  //----------------------------------------------------------------------------  
  void cmix_3track::hist_def() {
    // Define ntuples here
  }


  //----------------------------------------------------------------------------
  // begin_run() just contains some bookkeeping to make sure the numbers I save
  // are well-calibrated.
  //----------------------------------------------------------------------------    
  void cmix_3track::begin_run(BelleEvent* be, int* status) {
    BeamEnergy::begin_run();
    IpProfile::begin_run();
    eid::init_data();

    _runrow->expno = BeamEnergy::ExpNo();
    _runrow->runno = BeamEnergy::RunNo();
    
    _runrow->beam_e = BeamEnergy::E_beam_corr();
    _runrow->e_ler = BeamEnergy::E_LER();
    _runrow->e_her = BeamEnergy::E_HER();
    _runrow->cross_angle = BeamEnergy::Cross_angle();
    
    HepLorentzVector init = initial_state();

    _runrow->init_px = init.px();
    _runrow->init_py = init.py();
    _runrow->init_pz = init.pz();
    _runrow->init_e = init.e();

    Hep3Vector boost = init.boostVector();
    _runrow->boost_x = boost.x();
    _runrow->boost_y = boost.y();
    _runrow->boost_z = boost.z();        

    _runnt->append(*_runrow);
  } 

  //----------------------------------------------------------------------------
  // h5event() is called if the `h5' output format is requested from the
  // control script; again, it processes a single event and saves quantities of 
  // interest to the output file.
  //----------------------------------------------------------------------------  
  void cmix_3track::event(BelleEvent* evptr, int* status) {
    Belle_event_Manager& m_event = Belle_event_Manager::get_manager();

    _nevents_tot++;

    _countrow->expno   = m_event[0].ExpNo();
    _countrow->runno   = m_event[0].RunNo();
    _countrow->eventno = m_event[0].EvtNo();

    _countrow->kept   = 1;
    _countrow->npos  = -1;
    _countrow->nneg = -1;

    _countrow->npos_combos = -1;
    _countrow->nneg_combos = -1;
    
    
    Mdst_charged_Manager& m_charged = Mdst_charged_Manager::get_manager();
    _countrow->ncharged_raw = m_charged.count();
    
    if (IpProfile::usable()) {
      IP    = IpProfile::position(1);
      IPerr = IpProfile::position_err_b_life_smeared(1);    
    }
    else {
      IP    = HepPoint3D(0,0,0);
      IPerr = HepSymMatrix(3,0);
    }
            
    // Build charged track combos
    vector<Mdst_charged> vpos;
    vector<Mdst_charged> vneg;

    for (Mdst_charged_Manager::iterator it = m_charged.begin();
	 it != m_charged.end(); it++) 
      {
	if (!good_charged(*it))
	  continue;
	
	if (it->charge() > 0)
	  vpos.push_back(*it);
	else
	  vneg.push_back(*it);
      }
    _countrow->npos = vpos.size();
    _countrow->nneg = vneg.size();
                
    if ((_countrow->npos + _countrow->nneg) < 3) {
      _countrow->kept = -1;
      _countnt->append(*_countrow);
      return;
    }

    HepLorentzVector initial = initial_state(); 
    HepLorentzVector initial_cm = initial_state();

    Hep3Vector boost_cm      = initial.boostVector();   
    initial_cm.boost(-boost_cm);

    vector<vector<Mdst_charged> > trios;
    track_combinations(trios, vpos, vneg, 2, 1);
    _countrow->npos_combos = trios.size();
    track_combinations(trios, vpos, vneg, 1, 2);   
    _countrow->nneg_combos = trios.size() - _countrow->npos_combos;

    if (trios.size() == 0) {
      _countrow->kept = -2;
      _countnt->append(*_countrow);
      return;
    }

    _countnt->append(*_countrow);
    
    atc_pid pid_kpi(3,1,5,3,2);

    //    std::set<long> combo_hashes;

    pyt::table_writer<cmix_chargedrow_t>* _nt;
    cmix_chargedrow_t* _row; 
    for (vector<vector<Mdst_charged> >::iterator it = trios.begin(); 
	 it != trios.end(); it++) 
      {
	vector<Mdst_charged> trio = *it;
       
	double mag_p1 = sqrt(pow(trio[0].px(), 2) + pow(trio[0].py(), 2) + pow(trio[0].pz(), 2));
	double mag_p2 = sqrt(pow(trio[1].px(), 2) + pow(trio[1].py(), 2) + pow(trio[1].pz(), 2));

	if (mag_p1 < mag_p2) {
	  Mdst_charged tmp = trio[0];
	  trio[0] = trio[1];
	  trio[1] = tmp;
	}

	Mdst_charged slow = trio[1];

	Mdst_charged pos  = trio[0];
	Mdst_charged neg  = trio[2];
	
	if (slow.charge() < 0) {
	  Mdst_charged tmp = pos;
	  pos = neg;
	  neg = tmp;
	}			       	
		     
	HepLorentzVector p4pip = charged_track_p4(IP, pos, 2);
	HepLorentzVector p4kp  = charged_track_p4(IP, pos, 3);

	HepLorentzVector p4pim = charged_track_p4(IP, neg, 2);
	HepLorentzVector p4km  = charged_track_p4(IP, neg, 3);

	HepLorentzVector p4slow_pi = charged_track_p4(IP, slow, 2);
	HepLorentzVector p4slow_k  = charged_track_p4(IP, slow, 3);		       

	double m_kpi = (p4kp + p4pim).m();
	double m_pik = (p4km + p4pip).m();

	//	long idhash = three_body_idhash(slow.get_ID(), pos.get_ID(), neg.get_ID());
	  
	if (((m_kpi < 1.75) || (m_kpi > 2.0)) && ((m_pik < 1.75) || (m_pik > 2.0)))
	  continue; 

	double q_kpi = (p4slow_pi + p4kp + p4pim).m() - (p4kp + p4pim).m() - 0.13957018;
	double q_pik = (p4slow_pi + p4pip + p4km).m() - (p4pip + p4km).m() - 0.13957018;

	if ((q_kpi > 0.02) && (q_pik > 0.02))
	  continue;
       				
	// Decide which ntuple to write output to
	if (slow.charge() > 0) {
	  _row = _rowpos;
	  _nt  = _ntpos;
	}
	else {
	  _row = _rowneg;
	  _nt  = _ntneg;
	}

	_row->expno   = m_event[0].ExpNo();
	_row->runno   = m_event[0].RunNo();
	_row->eventno = m_event[0].EvtNo();
	
	//	cout << "SVD hitsrp " << trkFit.nhits(3) << endl;
	//cout << "SVD hitsz " << trkFit.nhits(4) << endl;

	_row->slow_svd_nhits_rphi  = slow.trk().mhyp(2).nhits(3);
	_row->slow_svd_nhits_z     = slow.trk().mhyp(2).nhits(4);
	_row->slow_svd_hitarr      = slow.trk().mhyp(2).hit_svd();	

	_row->idhash = three_body_idhash(slow.get_ID(), pos.get_ID(), neg.get_ID());	

	_row->pos_kpi_llh = pid_kpi.prob(pos);

	eid pos_eid(pos);
	_row->pos_eid = pos_eid.prob(3, -1, 5);

	_row->pip_svd_nhits_rphi  = pos.trk().mhyp(3).nhits(3);
	_row->pip_svd_nhits_z  = pos.trk().mhyp(3).nhits(4);
	_row->pip_svd_hitarr = pos.trk().mhyp(3).hit_svd();	

	_row->kp_svd_nhits_rphi  = pos.trk().mhyp(3).nhits(3);
	_row->kp_svd_nhits_z  = pos.trk().mhyp(3).nhits(4);
	_row->kp_svd_hitarr = pos.trk().mhyp(3).hit_svd();	

	_row->pip_px = p4pip.px();
	_row->pip_py = p4pip.py();
	_row->pip_pz = p4pip.pz();

	_row->kp_px = p4kp.px();
	_row->kp_py = p4kp.py();
	_row->kp_pz = p4kp.pz();

	_row->neg_kpi_llh = pid_kpi.prob(neg);

	eid neg_eid(neg);
	_row->neg_eid = neg_eid.prob(3, -1, 5);

	_row->pim_svd_nhits_rphi  = pos.trk().mhyp(3).nhits(3);
	_row->pim_svd_nhits_z  = pos.trk().mhyp(3).nhits(4);
	_row->pim_svd_hitarr = pos.trk().mhyp(3).hit_svd();	

	_row->km_svd_nhits_rphi  = pos.trk().mhyp(3).nhits(3);
	_row->km_svd_nhits_z  = pos.trk().mhyp(3).nhits(4);
	_row->km_svd_hitarr = pos.trk().mhyp(3).hit_svd();	

	_row->pim_px = p4pim.px();
	_row->pim_py = p4pim.py();
	_row->pim_pz = p4pim.pz();

	_row->km_px = p4km.px();
	_row->km_py = p4km.py();
	_row->km_pz = p4km.pz();
	
	_row->slow_kpi_llh = pid_kpi.prob(slow);

	_row->slow_px = slow.px();
	_row->slow_py = slow.py();
	_row->slow_pz = slow.pz();
	
	_row->m_kk   = (p4kp + p4km).m();
	_row->m_kpi  = (p4kp + p4pim).m();
	_row->m_pik  = (p4pip + p4km).m();
	_row->m_pipi = (p4pip + p4pim).m();
	
	_row->m_pipipi = (p4slow_pi+ p4pip + p4pim).m();
	_row->m_kkk    = (p4slow_k + p4kp + p4km).m();
	_row->m_pikpi  = (p4slow_pi + p4kp + p4pim).m();
	_row->m_pipik  = (p4slow_pi + p4pip + p4km).m();
	_row->m_kkpi   = (p4slow_k + p4kp + p4pim).m();

	_row->m_kpik   = (p4slow_k + p4pip + p4km).m();
	
	_row->q_kpi = (p4slow_pi + p4kp + p4pim).m() - (p4kp + p4pim).m() - 0.13957018;
	_row->q_pik = (p4slow_pi + p4pip + p4km).m() - (p4pip + p4km).m() - 0.13957018;

	Hep3Vector kpi_boost = (p4kp + p4pim).boostVector();
	HepLorentzVector p4kppi_rf = p4kp;
	HepLorentzVector p4kmpi_rf = p4pim;
	p4kppi_rf.boost(-kpi_boost);
	p4kmpi_rf.boost(-kpi_boost);

	_row->kpi_theta_star = kpi_boost.unit().dot(p4kmpi_rf.vect().unit());

	Hep3Vector pik_boost = (p4pip + p4km).boostVector();
	HepLorentzVector p4pipk_rf = p4pip;
	HepLorentzVector p4pimk_rf = p4km;
	p4pipk_rf.boost(-pik_boost);
	p4pimk_rf.boost(-pik_boost);

	_row->pik_theta_star = pik_boost.unit().dot(p4pimk_rf.vect().unit());

	HepLorentzVector p4cm_pikpi = p4slow_pi + p4kp + p4pim;
	HepLorentzVector p4cm_pipik = p4slow_pi + p4km + p4pip;		
	p4cm_pikpi.boost(-boost_cm);
	p4cm_pipik.boost(-boost_cm);
      
	_row->pcm_pikpi = p4cm_pikpi.vect().mag();
	_row->pcm_pipik = p4cm_pipik.vect().mag();
		
	_nt->append(*_row);
	_combos_found++;
      }
  }
  

  
  //----------------------------------------------------------------------------
  // end_run() handles any bookkeeping necessary at the end of an experimental
  // run--in this case it's empty.
  //----------------------------------------------------------------------------    
  void cmix_3track::end_run(BelleEvent* be, int* status) { 
    // Yessir
  }

  //----------------------------------------------------------------------------
  // term() is used to write out any output waiting to be flushed
  //----------------------------------------------------------------------------    
  void cmix_3track::term() {
    cout << "CRAIGB : " << _nevents_tot << " events processed, " <<  _combos_found << " D0/D0bar candidates found" << endl;    
    
    // End of analysis, close ntuples
  }


  //----------------------------------------------------------------------------
  // disp_stat() would display some statistics about the events I saved or threw
  // away if I was motivated enough to keep track of such things.
  //----------------------------------------------------------------------------    
  void cmix_3track::disp_stat(const char* str) {
    // Who knows
  }
                   
  //----------------------------------------------------------------------------
  // calc_R2() is a function that calculates a physics quantity (normalized Fox-
  // Wolfram second moment, if you're curious).
  //----------------------------------------------------------------------------    
    

#ifdef BELLE_NAMESPACE
}
#endif
